# ShaBa - Set of utilies for easy backups.
# Copyright (C) 2012 Damien Vincent
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

import argparse
import hashlib
import os.path
import os
import shutil
import signal
import sys

FILENAME_SHA1_SIZE = 40


stop_backup = 0
def sigint_handler(signal, frame):
  sys.stderr.write('Interrupting backup...\n')
  stop_backup = 1

# Calculate the hash of the content of a file
def compute_hash_file(f, hasher, blocksize=65536):
  buf = f.read(blocksize)
  while len(buf) > 0:
    hasher.update(buf)
    buf = f.read(blocksize)
  return hasher.digest()


# Check if the file is already in the backup directory.
# Here, only the content of the file matters
# (regardless of the file name)
def compare_file_hash(filename, directory_hash):
  try:
    f = file(filename, 'r')
  except IOError:
    return (0, [])
  f_hash = compute_hash_file(f, hashlib.sha1())
  hashname = ''.join( [ "%02x" % ord( x ) for x in f_hash ] )
  filename_hash = os.path.join(directory_hash, hashname)
  cmp_result = os.path.isfile(filename_hash)
  return (cmp_result, f_hash)


# Check if the file is already in the backup.
# Here, only the filename and the attributes matter.
# This is not as accurate as the hash version,
# but this is much more efficient
# def compare_file_attributes(filename, filename_ref)
# TODO


# Verify that a directory is empty
def is_empty_dir(dirname):
  for root, sub_folders, files in os.walk(dirname):
    if sub_folders or files:
      return 0
  return 1


# Verify that the directory is a directory that contains only hash files.
# => No subdirectories or other files.
def is_hash_dir(dirname):
  for root, sub_folders, files in os.walk(dirname):
    if sub_folders:
      return 0
    for f in files:
      if len(f) != FILENAME_SHA1_SIZE:
        return 0
  return 1


# Recursively calculate the total size of files inside |dirname|.
def backup_total_size(dirname):
  byte_count = 0
  for root, sub_folders, files in os.walk(dirname):
    for f in files:
      filename_source = os.path.join(root, f)
      if not os.path.islink(filename_source):
        byte_count += os.path.getsize(filename_source)
  return byte_count


# Do a backup of directory |backup_dir_source| to |backup_dir_target|
# The content of files are copied into backup_dir_hash.
def backup(backup_dir_source, backup_dir_target, backup_dir_hash, backup_size):
  byte_count = 0
  newfile_byte_count = 0

  for root, sub_folders, files in os.walk(backup_dir_source):
    root_relative = os.path.relpath(root, backup_dir_source)
    for d in sub_folders:
      try:
        os.mkdir(os.path.join(backup_dir_target, root_relative, d))
      except OSError:
        print ('Cannot create: ' + root_relative)
    for f in files:
      if stop_backup:
        return
      filename_relative = os.path.join(root_relative, f)
      progress = (100 * byte_count) / backup_size
      filename_display = filename_relative + (" " * 80)
      sys.stderr.write((('\r[%3d%%]: ' % progress) + filename_display[0:79]))
      filename_source = os.path.join(backup_dir_source, root_relative, f)
      filename_target = os.path.join(backup_dir_target, root_relative, f)
      if os.path.islink(filename_source):
        # Populate the backup directory
        filename_symlink = os.readlink(filename_source)
        os.symlink(filename_symlink, filename_target)
        # TODO: attributes of the sym links
      else:
        # Populate the hash directory (if needed)
        (r, f_hash) = compare_file_hash(filename_source, backup_dir_hash)
        hashname = ''.join( [ "%02x" % ord( x ) for x in f_hash ] )
        filename_hash = os.path.join(backup_dir_hash, hashname)
        if r == 0:
          newfile_byte_count += os.path.getsize(filename_source)
          shutil.copy2(filename_source, filename_hash)
        # Populate the backup directory
        os.link(filename_hash, filename_target)
        # Update statistics
        byte_count += os.path.getsize(filename_source)
  sys.stderr.write('\n')
  sys.stderr.write('Total new files: ' + ('%d' % newfile_byte_count) + ' bytes\n')

  # Set the same permissions & modif time
  for root, sub_folders, files in os.walk(backup_dir_source):
    root_relative = os.path.relpath(root, backup_dir_source)
    for d in sub_folders:
      shutil.copystat(os.path.join(backup_dir_source, root_relative, d),
                      os.path.join(backup_dir_target, root_relative, d))


def main():
  signal.signal(signal.SIGINT, sigint_handler)

  parser = argparse.ArgumentParser(description='Backup a directory.')
  parser.add_argument('--src', action="store", dest="backup_dir_source",
                      help='Directory to backup')
  parser.add_argument('--dst', action="store", dest="backup_dir_target",
                      help='Backup directory')
  parser.add_argument('--hash', action="store", dest="backup_dir_hash",
                      help='Directory to store the hash')
  args = parser.parse_args()

  backup_dir_source = args.backup_dir_source
  backup_dir_hash   = args.backup_dir_hash
  backup_dir_target = args.backup_dir_target

  sys.stderr.write('Step #1: Checking directories\n')
  if not os.path.exists(backup_dir_source):
    sys.stderr.write('Error: Source directory does not exist\n')
  if not os.path.exists(backup_dir_target):
    sys.stderr.write('Error: Target directory does not exist\n')
    sys.exit(1)
  if not is_empty_dir(backup_dir_target):
    sys.stderr.write('Error: Target directory must be empty\n')
    sys.exit(1)
  if not is_hash_dir(backup_dir_hash):
    sys.stderr.write('Error: Hash directory contains some sub directories or common files\n')
    sys.exit(1)

  sys.stderr.write('Step #2: Collecting backup statistics\n')
  backup_size = backup_total_size(backup_dir_source)
  sys.stderr.write(('Total size to backup: %d bytes\n' % backup_size))

  sys.stderr. write('Step #3: Backup\n')
  backup(backup_dir_source, backup_dir_target, backup_dir_hash, backup_size)

if __name__ == "__main__":
  main()
